# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time : 2020/7/8 10:48 
@Author : SPZ
@File : customize_service.py 
@Software: pycharm
"""

import os
import threading

import numpy as np
import tensorflow as tf
from PIL import Image
from predict_lib.preprocess import process_imgs
from predict_lib.decode import decode

from model_service.tfserving_model_service import TfServingBaseService
import logging

logger = logging.getLogger(__name__)
ch = os.path.dirname(__file__)
dict_file_path = os.path.join(ch, 'digitEn_37.txt')
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


class OcrService(TfServingBaseService):

    def _preprocess(self, data):
        # http 两种请求形式
        # form-data 文件格式的请求对应 data = {"请求key值":{"文件名":<文件io>}}
        # json格式对应 data = json.loads("接口传入的json体")
        preprocessed_data = {}

        for k, v in data.items():
            for file_name, file_content in v.items():
                image1 = Image.open(file_content)
                input_tensors = process_imgs([image1])
                preprocessed_data[k] = input_tensors

        return preprocessed_data

    def _postprocess(self, data):
        infer_output = {"result": []}
        for output_name, results in data.items():
            infer_output["result"].append(decode(results, dict_file_path=dict_file_path))

        return infer_output
